#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <stdbool.h>
#include <ctype.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <arpa/inet.h>

#define DEF_DNS_ADDRESS "114.114.114.114"
#define LOCAL_ADDRESS "127.0.0.1"
#define DNS_PORT 53
#define BUF_SIZE 512
#define LENGTH 65
#define AMOUNT 300
#define NOTFOUND -1
#define TRUE 1
#define FALSE 0


//DNS header
struct DNSHeader
{
	unsigned short ID;
	unsigned short Flags;
	unsigned short QuestNum;
	unsigned short AnswerNum;
	unsigned short AuthorNum;
	unsigned short AdditionNum;
};
typedef struct DNSHeader  DNSHDR,*pDNSHDR;

//DNS translation table
struct translate
{
	char IP[24];  //IP address
	char domain[40]; // domain name
};

struct IDChange
{
	unsigned short oldID; //Original ID
	bool done; //mark if it done
	struct sockaddr_in client;
};
typedef struct IDChange IDTransform;

struct translate DNS_table[AMOUNT];
IDTransform IDTransTable[AMOUNT];
int IDcount = 0;
char url[LENGTH];

//去除字符串两边的空格
void strim(char *s)
{
        int end = 0;
        char *p=s;
        if(p)
        {
                while(*p&&isspace(*p)) p++;
                for(end = strlen(p)-1;isspace(p[end]);end--) ;
                p[end+1]='\0';
                strcpy(s,p);
        }
}

//构造DNS 解析表
int GetTable(char *tablePath)
{
	FILE *fp;
	char buff[512];
        char *item;
	char ip[24];
	char domain[40];
	int i=0;
	if ((fp=fopen(tablePath,"r")) == NULL)
	{
		printf("Cannot open file,make sure the file exists please\n");
		exit(1);
	}	
	while(fgets(buff,512,fp)) {
                item = strtok(buff," ");
		strcpy(DNS_table[i].IP,item);	
		strim(DNS_table[i].IP);
                item = strtok(NULL,"\n");
		strcpy(DNS_table[i].domain,item);
		strim(DNS_table[i].domain);
		i++;
        }
	fclose(fp);
	return i;
}
void GetUrl(char *recvbuf,int recvnum)
{
	char urlname[LENGTH];
	int i = 0, j, k = 0;

	memset(url, 0, LENGTH);
	memcpy(urlname, &(recvbuf[sizeof(DNSHDR)]), recvnum-16);

	int len = strlen(urlname);
	
	//域名转换
	while (i < len) {
		if (urlname[i] > 0 && urlname[i] <= 63)
			for (j = urlname[i], i++; j > 0; j--, i++, k++)
				url[k] = urlname[i];
		
		if (urlname[i] != 0) {
			url[k] = '.';
		    k++;
		}
	}

	url[k] = '\0';
}

int IsFind(char* url, int num)
{
	int find = NOTFOUND;
	char* domain;
	int i=0;

	for (i = 0; i < num; i++) {
		domain = (char *)DNS_table[i].domain;
		if (strcmp(domain, url) == 0) {	//找到
			find = i;
			break;
		}
	}

	return find;
}
unsigned short RegisterNewID (unsigned short oID, struct sockaddr_in temp, bool ifdone)
{
	srand(time(NULL));
	IDTransTable[IDcount].oldID = oID;
	IDTransTable[IDcount].client = temp;
	IDTransTable[IDcount].done  = ifdone;
	IDcount++;

	return (unsigned short)(IDcount-1);	//以表中下标作为新的ID
}

int main(int argc,char ** argv)
{
	int socketServer,socketLocal;
	struct sockaddr_in serverName,clientName,localName;
	char sendbuf[BUF_SIZE];
	char recvbuf[BUF_SIZE];
	char tablePath[100];
	char outerDns[16];
	int iLen_cli,iSend,iRecv;
	int num;
	
	if(argc == 1){
		strcpy(outerDns,DEF_DNS_ADDRESS);
		strcpy(tablePath, "mytest.txt");
	}
	else if(argc == 2) {
		strcpy(outerDns, argv[1]);
		strcpy(tablePath, "mytest.txt");
	}
	else if(argc == 3) {
		strcpy(outerDns,argv[1]);
		strcpy(tablePath,argv[2]);
	}
	
	num = GetTable(tablePath);
	for (int i=0; i < AMOUNT; i++) {  //初始化ID转换表
		IDTransTable[i].oldID = 0;
		IDTransTable[i].done  = 0;
		memset(&(IDTransTable[i].client), 0, sizeof(struct sockaddr_in));
	}
	
	//创建本地DNS和外部DNS套接字
    	socketServer = socket(AF_INET, SOCK_DGRAM, 0);
	socketLocal = socket(AF_INET, SOCK_DGRAM, 0);

	//设置本地DNS和外部DNS两个套接字
	localName.sin_family = AF_INET;
	localName.sin_port = htons(DNS_PORT);
	localName.sin_addr.s_addr = inet_addr(LOCAL_ADDRESS);

	serverName.sin_family = AF_INET;
	serverName.sin_port = htons(DNS_PORT);
	serverName.sin_addr.s_addr = inet_addr(outerDns);
	
	//绑定本地DNS服务器地址
	if (bind(socketLocal, (struct sockaddr*)&localName, sizeof(localName))) {
		printf("Binding Port 53 failed.\n");
		exit(1);
	}
	else
		printf("Binding Port 53 succeed.\n"); 

	//本地DNS中继服务器的具体操作
	while (1) {
		iLen_cli = sizeof(clientName);
	    memset(recvbuf, 0, BUF_SIZE);

		//接收DNS请求
		iRecv = recvfrom(socketLocal, recvbuf, sizeof(recvbuf), 0, (struct sockaddr*)&clientName, &iLen_cli);

		if (iRecv == -1) {
			printf("Recvfrom Failed:\n");
			continue;
		}
		else if (iRecv == 0) {
			break;
		}
		else {	
			GetUrl(recvbuf, iRecv);				//获取域名
			int find = IsFind(url, num);		//在域名解析表中查找

			printf("%s\n",url);

			//在域名解析表中没有找到
			if (find == NOTFOUND) {
				//ID转换
				unsigned short *pID = (unsigned short *)malloc(sizeof(unsigned short));
				memcpy(pID, recvbuf, sizeof(unsigned short));
				unsigned short nID = htons(RegisterNewID(ntohs(*pID), clientName, FALSE));
				memcpy(recvbuf, &nID, sizeof(unsigned short));

				//打印 时间 newID 功能 域名 IP
				//DisplayInfo(ntohs(nID), find);

				//把recvbuf转发至指定的外部DNS服务器
				iSend = sendto(socketServer, recvbuf, iRecv, 0, (struct sockaddr*)&serverName, sizeof(serverName));
				if (iSend == -1) {
					printf("sendto Failed\n");
					continue;
				}
				else if (iSend == 0)
					break;

				free(pID);	//释放动态分配的内存

				//接收来自外部DNS服务器的响应报文
				iRecv = recvfrom(socketServer, recvbuf, sizeof(recvbuf), 0, (struct sockaddr*)&clientName, &iLen_cli);

				//ID转换
				pID = (unsigned short *)malloc(sizeof(unsigned short));
				memcpy(pID, recvbuf, sizeof(unsigned short));
				int m = ntohs(*pID);
				unsigned short oID = htons(IDTransTable[m].oldID);
				memcpy(recvbuf, &oID, sizeof(unsigned short));
				IDTransTable[m].done = TRUE;

				//从ID转换表中获取发出DNS请求者的信息
				clientName = IDTransTable[m].client;

				//把recvbuf转发至请求者处
				iSend = sendto(socketLocal, recvbuf, iRecv, 0, (struct sockaddr*)&clientName, sizeof(clientName));
				if (iSend == -1) {
					printf("sendto Failed\n");
					continue;
				}
				else if (iSend == 0)
					break;

				free(pID);	//释放动态分配的内存
			}

			//在域名解析表中找到
			else {	
				//获取请求报文的ID
				unsigned short *pID = (unsigned short *)malloc(sizeof(unsigned short));
				memcpy(pID, recvbuf, sizeof(unsigned short));

				//转换ID
				unsigned short nID = RegisterNewID(ntohs(*pID), clientName, FALSE);

				//打印 时间 newID 功能 域名 IP
				//DisplayInfo(nID, find);

				//构造响应报文返回
				memcpy(sendbuf, recvbuf, iRecv);						//拷贝请求报文
				unsigned short a = htons(0x8180);
				memcpy(&sendbuf[2], &a, sizeof(unsigned short));		//修改标志域

				//修改回答数域
				if (strcmp(DNS_table[find].IP, "0.0.0.0") == 0)	
					a = htons(0x0000);	//屏蔽功能：回答数为0
				else
					a = htons(0x0001);	//服务器功能：回答数为1
				memcpy(&sendbuf[6], &a, sizeof(unsigned short));
				int curLen = 0;

				//构造DNS响应部分
				char answer[16];
				unsigned short Name = htons(0xc00c);
				memcpy(answer, &Name, sizeof(unsigned short));
				curLen += sizeof(unsigned short);

				unsigned short TypeA = htons(0x0001);
				memcpy(answer+curLen, &TypeA, sizeof(unsigned short));
				curLen += sizeof(unsigned short);

				unsigned short ClassA = htons(0x0001);
				memcpy(answer+curLen, &ClassA, sizeof(unsigned short));
				curLen += sizeof(unsigned short);

				unsigned long timeLive = htonl(0x7b);
				memcpy(answer+curLen, &timeLive, sizeof(unsigned long));
				curLen += sizeof(unsigned long);

				unsigned short IPLen = htons(0x0004);
				memcpy(answer+curLen, &IPLen, sizeof(unsigned short));
				curLen += sizeof(unsigned short);

				unsigned long IP = (unsigned long) inet_addr(DNS_table[find].IP);
				memcpy(answer+curLen, &IP, sizeof(unsigned long));
				curLen += sizeof(unsigned long);
				curLen += iRecv;

				//请求报文和响应部分共同组成DNS响应报文存入sendbuf
				memcpy(sendbuf+iRecv, answer, curLen);

				//发送DNS响应报文
				iSend = sendto(socketLocal, sendbuf, curLen, 0, (struct sockaddr*)&clientName, sizeof(clientName));
				if (iSend == -1) {
					printf("sendto Failed\n");
					continue;
				}
				else if (iSend == 0)
					break;
			
				free(pID);		//释放动态分配的内存
			}
		}
	}
}
